#include <thrust/device_ptr.h>
#include <thrust/sort.h>
#include <thrust/reduce.h>
#include <thrust/fill.h>
#include <thrust/iterator/constant_iterator.h>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <cuda_runtime.h>
#include <cub/cub.cuh>
#include <iostream>

void cubHistogram(const unsigned int* d_vals,
                     unsigned int* d_histo,
                     unsigned int numBins,
                     unsigned int numElems)
{
    // Step 1: 查询临时存储大小
    void* d_temp_storage = nullptr;
    size_t temp_storage_bytes = 0;

    // 注意：num_levels = numBins + 1
    cub::DeviceHistogram::HistogramEven(
        d_temp_storage,           // [in] temp storage
        temp_storage_bytes,       // [in|out] temp storage size
        d_vals,                   // [in] input samples
        d_histo,                  // [out] histogram (length = numBins)
        numBins + 1,              // [in] num_levels → bins = num_levels - 1
        static_cast<unsigned int>(0),        // [in] lower_level (inclusive)
        static_cast<unsigned int>(numBins),  // [in] upper_level (exclusive)
        numElems                 
    );

    // Step 2: 分配临时存储
    cudaMalloc(&d_temp_storage, temp_storage_bytes);

    if (!d_temp_storage) {
        fprintf(stderr, "Failed to allocate temporary storage\n");
        return;
    }

    // Step 3: 执行直方图计算
    cub::DeviceHistogram::HistogramEven(
        d_temp_storage,
        temp_storage_bytes,
        d_vals,
        d_histo,
        numBins + 1,
        static_cast<unsigned int>(0),
        static_cast<unsigned int>(numBins),
        numElems
    );
    // Step 4: 清理
    cudaFree(d_temp_storage);
}

void thrustHistogram(const unsigned int* d_vals, 
                     unsigned int* d_histo, 
                     unsigned int numBins, 
                     unsigned int numElems) 
{
    // 包装设备指针
    thrust::device_ptr<const unsigned int> d_vals_ptr(d_vals);
    thrust::device_ptr<unsigned int>       d_histo_ptr(d_histo);

    // Step 1: 拷贝并排序
    thrust::device_vector<unsigned int> sorted(d_vals, d_vals + numElems);
    thrust::sort(sorted.begin(), sorted.end());

    // Step 2: reduce_by_key
    thrust::device_vector<unsigned int> keys(numElems);
    thrust::device_vector<unsigned int> counts(numElems);

    auto new_end = thrust::reduce_by_key(
        sorted.begin(), sorted.end(),
        thrust::constant_iterator<unsigned int>(1),
        keys.begin(),
        counts.begin()
    );

    int numUnique = new_end.first - keys.begin();

    // Step 3: 初始化直方图（在设备上）
    thrust::fill(d_histo_ptr, d_histo_ptr + numBins, 0);

    // Step 4: 将 keys 和 counts 拷贝到主机进行处理
    thrust::host_vector<unsigned int> h_keys = keys;
    thrust::host_vector<unsigned int> h_counts = counts;

    // 现在在主机上操作：但我们要把结果写回设备指针 d_histo
    // 方法：使用 thrust::device_ptr 写入
    for (int i = 0; i < numUnique; ++i) {
        unsigned int key = h_keys[i];
        if (key < numBins) {
            d_histo_ptr[key] = h_counts[i];
        } 
    }

    cudaDeviceSynchronize();
    cudaGetLastError();
}

int main() {
    const unsigned int numElems = 10;
    const unsigned int numBins = 5;

    // 主机数据
    unsigned int h_vals[] = {0, 1, 1, 2, 2, 2, 3, 3, 4, 0}; // 值范围 [0,4]
    unsigned int *d_vals;
    unsigned int *d_histo;

    // 分配设备内存
    cudaMalloc(&d_vals, numElems * sizeof(unsigned int));
    cudaMalloc(&d_histo, numBins * sizeof(unsigned int));

    // 拷贝输入数据到设备
    cudaMemcpy(d_vals, h_vals, numElems * sizeof(unsigned int), cudaMemcpyHostToDevice);

    // 创建 CUDA 事件用于计时
    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);

    float time_thrust = 0.0f;
    float time_cub = 0.0f;

    // ----------------------------
    // 测试 thrustHistogram
    // ----------------------------
    std::cout << "\n[TEST] Running thrustHistogram..." << std::endl;

    // 重置直方图内存
    cudaMemset(d_histo, 0, numBins * sizeof(unsigned int));

    cudaEventRecord(start);
    thrustHistogram(d_vals, d_histo, numBins, numElems);
    cudaEventRecord(stop);

    cudaEventSynchronize(stop);
    cudaEventElapsedTime(&time_thrust, start, stop);

    // 拷贝结果回主机并打印
    unsigned int h_histo_thrust[numBins];
    cudaMemcpy(h_histo_thrust, d_histo, numBins * sizeof(unsigned int), cudaMemcpyDeviceToHost);

    std::cout << "[RESULT] Thrust Histogram:" << std::endl;
    for (int i = 0; i < numBins; ++i) {
        std::cout << "Bin[" << i << "] = " << h_histo_thrust[i] << std::endl;
    }


    // ----------------------------
    // 测试 cubHistogram
    // ----------------------------
    std::cout << "\n[TEST] Running cubHistogram..." << std::endl;

    // 重置直方图内存
    cudaMemset(d_histo, 0, numBins * sizeof(unsigned int));

    cudaEventRecord(start);
    cubHistogram(d_vals, d_histo, numBins, numElems);
    cudaEventRecord(stop);

    cudaEventSynchronize(stop);
    cudaEventElapsedTime(&time_cub, start, stop);

    // 拷贝结果回主机并打印
    unsigned int h_histo_cub[numBins];
    cudaMemcpy(h_histo_cub, d_histo, numBins * sizeof(unsigned int), cudaMemcpyDeviceToHost);

    std::cout << "[RESULT] CUB Histogram:" << std::endl;
    for (int i = 0; i < numBins; ++i) {
        std::cout << "Bin[" << i << "] = " << h_histo_cub[i] << std::endl;
    }


    // ----------------------------
    // 输出时间对比
    // ----------------------------
    std::cout << "\n[PERFORMANCE] Timing Results:" << std::endl;
    std::cout << "  Thrust Histogram Time: " << time_thrust << " ms" << std::endl;
    std::cout << "  CUB Histogram Time:    " << time_cub    << " ms" << std::endl;
    std::cout << "  Speedup (Thrust/CUB):  " << (time_cub > 0 ? time_thrust / time_cub : 0.0f) << "x" << std::endl;

    // 清理
    cudaFree(d_vals);
    cudaFree(d_histo);
    cudaEventDestroy(start);
    cudaEventDestroy(stop);

    return 0;
}